from fastapi import APIRouter, Request, WebSocket, WebSocketDisconnect
from pydantic import ValidationError
from common.utils import FileUtil, generate_unique_str
from common.log import LOG, LogLevel
from runapp.func import make_script, run_script, watch_for_stop, read_and_send_stderr, read_and_send_stdout
from starlette.websockets import WebSocketState
from runapp.schemas import WsRequest
import asyncio
import signal


router = APIRouter(prefix="/run", tags=["运行模块"])


temp_path = './runapp/temp/' # 临时文件夹的路径


@router.get('/testfile', description="测试文件类是否可用")
async def test():
    filename = temp_path + generate_unique_str() + ".temp"
    await FileUtil(filename).create_file()
    await FileUtil(filename).write_file("print('hello world')" + "\n")
    await FileUtil(filename).write_file("print('hello world')" + "\n")
    print(await FileUtil(filename).read_file())
    await FileUtil(filename).delete_file()


@router.websocket('/run')
async def startrun(websocket: WebSocket):
    """
    1. 客户端connect，还什么都没干，就直接disconnect
	2. 客户端connect后，发送无效数据，被服务端主动disconnect
	3. 客户端connect并发送有效数据后，启动了循环监控脚本，之后客户端发生异常，主动disconnect
	4. 客户端connect并发送有效数据后，启动了循环监控脚本，之后监控脚本发生异常，客户端收到该异常，发送stop信号终止包括监控脚本在内的一系列任务
    5. 客户端connect并发送有效数据后，启动了循环监控脚本，之后主动发送stop信号，终止任务。
    """
    await websocket.accept()
    stop_listener = None
    proc = None
    stdout_listener = None
    stderr_listener = None

    try:
        # 1. 接收并验证初始数据（只接收一次）
        raw_data = await websocket.receive_json()
        try:
            request = WsRequest(**raw_data)
        except ValidationError as e:
            await websocket.send_json({
                "status": "error",
                "type": "INVALID_FORMAT",
                "details": "数据格式有误"
            })
            LOG(LogLevel.INFO) << "客户端传入无效数据！" << "\n"
            return  # 格式错误直接终止连接

        # 2. 创建任务取消标志
        stop_event = asyncio.Event()

        # 3. 启动stop监听任务
        stop_listener = asyncio.create_task(watch_for_stop(websocket, stop_event))

        # 启动监控脚本
        proc = await asyncio.create_subprocess_exec(
            "python", f"{temp_path}test.py",
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            bufsize=0
        )

        # 创建两个任务并行读取 stdout 和 stderr
        stdout_listener = asyncio.create_task(read_and_send_stdout(websocket, proc.stdout, "STDOUT"))
        stderr_listener = asyncio.create_task(read_and_send_stderr(websocket, proc.stderr, "STDERR"))
            
        while True:
            """
            如果主循环 (while True) 中没有 await 语句
            这会导致循环变成紧密循环，100% 占用 CPU 且不释放控制权
            结果：事件循环被阻塞，无法处理其他任务（包括 stop_listener 任务）
            """
            if websocket.client_state == WebSocketState.DISCONNECTED:
                raise

            if stop_event.is_set():
                LOG(LogLevel.INFO) << "接收到stop命令，终止实时监控任务！" << "\n" 
                break
            
            await asyncio.sleep(0.1) # 必须有这行代码！！！
        

    except (WebSocketDisconnect, RuntimeError):
        LOG(LogLevel.INFO) << "客户端异常，断开连接！" << "\n"
    except Exception as e:
        LOG(LogLevel.ERROR) << f"发生了未知异常: {str(e)}" << "\n"
        await websocket.send_json({
            "status": "error",
            "type": "PROCESSING_ERROR",
            "details": str(e)
        })
    finally:
        # 安全清理
        if stop_listener is not None:  # 更严格的None检查
            if not stop_listener.done():  # 只有未完成的任务需要取消
                stop_listener.cancel()
                try:
                    await asyncio.wait_for(stop_listener, timeout=2.0)  # 带超时等待
                except asyncio.CancelledError:
                    LOG(LogLevel.DEBUG) << "stop监听任务正常取消" << "\n"
                except asyncio.TimeoutError:
                    LOG(LogLevel.WARNING) << "stop监听任务取消超时，强制结束" << "\n"
                except Exception as e:
                    LOG(LogLevel.ERROR) << f"stop监听任务取消异常: {type(e).__name__}:{e}" << "\n"
            else:
                LOG(LogLevel.DEBUG) << "stop监听任务完成（done），资源已释放！" << "\n"
      
        if proc is not None:
            if proc.returncode is None:
                # 关闭流，进而触发读取stdout和stderr的协程退出
                if proc.stdout:
                    proc.stdout._transport.close()  # 强制关闭底层传输（丢弃后续数据）
                if proc.stderr:
                    proc.stderr._transport.close() 
                    
                # 如果监控脚本还在运行，就发送 ctrl + c 信号让其终止
                LOG(LogLevel.INFO) << "向监控脚本发送 Ctrl+C 信号！" << "\n"
                proc.send_signal(signal.SIGINT)

            await proc.wait()
            LOG(LogLevel.DEBUG) << "进程资源已释放！" << "\n"
            LOG(LogLevel.INFO) << f"监控脚本已退出，退出码: {proc.returncode}" << "\n"
            await stdout_listener
            LOG(LogLevel.DEBUG) << "stdout监控协程资源已释放！" << "\n"
            await stderr_listener
            LOG(LogLevel.DEBUG) << "stderr监控协程资源已释放！" << "\n"

        if websocket.client_state != WebSocketState.DISCONNECTED:
            await websocket.close()

        